# eval_embeddings.py
# usages: 评估embeddings模型出来的结果是否靠谱
# python ./src/tools/eval_embeddings.py --samples "./data/samples.jsonl" > tmp_run_eval_embeddings.log

import os,sys
import argparse
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)
import json
import sys
from pathlib import Path
from query.query_vector import retrieve_similar_texts
from query.extract_tables import extract_tables


def load_samples(file_path):
    """读取 JSON Lines 格式的样本文件"""
    samples = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                sample = json.loads(line)
                # 验证必要字段
                if 'id' not in sample or 'question' not in sample or 'sql' not in sample:
                    print(f"警告：第 {line_num} 行缺少必要字段: {line}", file=sys.stderr)
                    continue
                samples.append(sample)
            except json.JSONDecodeError as e:
                print(f"第 {line_num} 行 JSON 解析错误: {e}", file=sys.stderr)
    return samples


def main():
    parser = argparse.ArgumentParser(description="评估 embeddings模型: SQL 和 DDL 提取的表是否一致")
    parser.add_argument('--samples', type=str, required=True, help='输入样本文件路径 (JSON Lines 格式)')
    args = parser.parse_args()

    samples_file = Path(args.samples)
    if not samples_file.exists():
        print(f"错误：文件不存在: {samples_file}", file=sys.stderr)
        sys.exit(1)

    samples = load_samples(samples_file)

    print(f"共加载 {len(samples)} 条样本", file=sys.stderr)

    # 打开日志文件，用于记录不匹配项和最终结果
    log_file = "tmp_eval_embeddings.log"
    mismatch_count = 0

    with open(log_file, 'w', encoding='utf-8') as log_f:
        for sample in samples:
            try:
                question = sample['question']
                sql = sample.get('sql', '')
                sample_id = sample['id']
                level = sample.get('level', 'N/A')

                # 步骤1: 获取 DDL
                filtered_docs = retrieve_similar_texts(question, 1,  top_k=10)
                try:
                    matched_doc = json.loads(filtered_docs[0])
                    ddl = matched_doc.get("DDL", "")
                except Exception as e:
                    print(f"❌ 解析 JSON 失败: {e}", file=sys.stderr)
                    ddl = ""

                # 步骤2: 提取 SQL 和 DDL 中的表名
                sql_tables = extract_tables(sql)
                ddl_tables = extract_tables(ddl)

                # 转为集合进行比较（不区分顺序）
                if set(sql_tables) != set(ddl_tables):
                    mismatch_count += 1

                    # 写入日志文件
                    log_f.write("-" * 60 + "\n")
                    log_f.write(f"ID: {sample_id}\n")
                    log_f.write(f"Level: {level}\n")
                    log_f.write(f"Question: {question}\n")
                    log_f.write(f"SQL: {sql}\n")
                    log_f.write(f"DDL: {ddl}\n")
                    log_f.write(f"SQL 表: {sorted(sql_tables)}\n")
                    log_f.write(f"DDL 表: {sorted(ddl_tables)}\n")
                    log_f.flush()  # 立即写入磁盘，便于实时查看

                    # 同时在 stderr 输出提示（可选）
                    print(f"不匹配样本 ID={sample_id}:{question}:sql-{sql_tables}:ddl-{ddl_tables} 已记录", file=sys.stderr)

            except Exception as e:
                print(f"处理样本 ID={sample.get('id', 'unknown')} 时发生异常: {e}", file=sys.stderr)
                continue

        # 写入最终统计到日志文件
        final_report = f"{'-' * 60}\n检查完成。不匹配数量: {mismatch_count}/{len(samples)}\n"
        log_f.write(final_report)
        log_f.flush()

        # 同时输出到 stderr
        print(final_report.strip(), file=sys.stderr)

    print(f"结果已写入 {log_file}", file=sys.stderr)

if __name__ == '__main__':
    main()